#!/bin/bash
# Z-Image-Turbo NPU 部署脚本
# 模型路径: /data0/Z-Image-Turbo

set -e

# 配置
MODEL_PATH="/data0/Z-Image-Turbo"
CONTAINER_NAME="z-image-turbo-api"
IMAGE_NAME="z-image-turbo-npu:latest"
HOST_PORT=18004
CONTAINER_PORT=8000
NPU_DEVICE="davinci3"  # 可根据需要修改: davinci0, davinci1, davinci2, davinci3

echo "=========================================="
echo "Z-Image-Turbo NPU 部署脚本"
echo "=========================================="

# 检查Ascend CANN环境
echo "检查Ascend CANN环境..."
ASCEND_DRIVER="/usr/local/Ascend/driver"
ASCEND_DRIVER_LIB="/usr/local/Ascend/driver/lib64"
ASCEND_DRIVER_VERSION="/usr/local/Ascend/driver/version.info"
ASCEND_TOOLKIT="/usr/local/Ascend/ascend-toolkit"
ASCEND_INSTALL_INFO="/etc/ascend_install.info"
DCMI_PATH="/usr/local/dcmi"
NPU_SMI="/usr/local/bin/npu-smi"

# 检查NPU设备
echo "检查NPU设备..."
npu-smi info || echo "警告: 无法执行npu-smi，请确保已安装昇腾驱动"

# 检查关键路径和文件
MISSING_PATHS=()

if [ ! -d "$ASCEND_DRIVER" ]; then
    echo "警告: 未找到 $ASCEND_DRIVER 目录"
    MISSING_PATHS+=("$ASCEND_DRIVER")
else
    echo "✓ 找到驱动目录: $ASCEND_DRIVER"
    if [ -d "$ASCEND_DRIVER_LIB" ]; then
        if find "$ASCEND_DRIVER_LIB" -name "libascend_hal.so" 2>/dev/null | head -1 | grep -q .; then
            echo "✓ 找到 libascend_hal.so"
        else
            echo "警告: 在 $ASCEND_DRIVER_LIB 中未找到 libascend_hal.so"
        fi
    fi
fi

if [ ! -d "$ASCEND_TOOLKIT" ]; then
    echo "提示: 未找到 $ASCEND_TOOLKIT 目录（可选，但推荐挂载）"
else
    echo "✓ 找到CANN工具包目录: $ASCEND_TOOLKIT"
fi

if [ ! -f "$ASCEND_DRIVER_VERSION" ]; then
    echo "警告: 未找到 $ASCEND_DRIVER_VERSION 文件"
    MISSING_PATHS+=("$ASCEND_DRIVER_VERSION")
else
    echo "✓ 找到驱动版本信息: $ASCEND_DRIVER_VERSION"
fi

if [ ! -f "$ASCEND_INSTALL_INFO" ]; then
    echo "警告: 未找到 $ASCEND_INSTALL_INFO 文件"
    MISSING_PATHS+=("$ASCEND_INSTALL_INFO")
else
    echo "✓ 找到安装信息: $ASCEND_INSTALL_INFO"
fi

if [ ! -d "$DCMI_PATH" ]; then
    echo "提示: 未找到 $DCMI_PATH 目录（可选）"
else
    echo "✓ 找到DCMI路径: $DCMI_PATH"
fi

if [ ! -f "$NPU_SMI" ]; then
    echo "提示: 未找到 $NPU_SMI 工具（可选）"
else
    echo "✓ 找到 npu-smi 工具: $NPU_SMI"
fi

if [ ${#MISSING_PATHS[@]} -gt 0 ]; then
    echo ""
    echo "警告: 以下路径未找到，容器可能无法正常访问NPU运行时库："
    for path in "${MISSING_PATHS[@]}"; do
        echo "  - $path"
    done
    echo "这可能导致 libascend_hal.so 错误"
    echo ""
fi

# 检查模型目录
if [ ! -d "$MODEL_PATH" ]; then
    echo "错误: 模型目录 $MODEL_PATH 不存在"
    echo "请先下载模型:"
    echo "  huggingface-cli download Tongyi-MAI/Z-Image-Turbo --local-dir $MODEL_PATH"
    echo ""
    echo "或者使用以下命令（需要设置 HF_XET_HIGH_PERFORMANCE=1 以使用 XET 下载）:"
    echo "  HF_XET_HIGH_PERFORMANCE=1 huggingface-cli download Tongyi-MAI/Z-Image-Turbo --local-dir $MODEL_PATH"
    exit 1
else
    echo "✓ 找到模型目录: $MODEL_PATH"
    # 检查关键文件
    if [ -f "$MODEL_PATH/model_index.json" ] || [ -f "$MODEL_PATH/config.json" ]; then
        echo "✓ 模型文件检查通过"
    else
        echo "警告: 模型目录存在但可能不完整，请确认模型已正确下载"
    fi
fi

# 构建Docker镜像
echo ""
echo "构建Docker镜像..."
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
cd "${SCRIPT_DIR}/z-image-turbo"
docker build -t ${IMAGE_NAME} .
cd "${SCRIPT_DIR}"

# 停止并删除旧容器
echo ""
echo "停止旧容器..."
docker stop ${CONTAINER_NAME} 2>/dev/null || true
docker rm ${CONTAINER_NAME} 2>/dev/null || true

# 启动新容器
echo ""
echo "启动 Z-Image-Turbo 服务 (NPU: ${NPU_DEVICE}, 端口: ${HOST_PORT})..."

# 从NPU设备名中提取设备ID（例如：davinci3 -> 3）
NPU_DEVICE_ID=$(echo ${NPU_DEVICE} | sed 's/davinci//')

# 检查指定的NPU设备是否存在
NPU_DEVICE_PATH="/dev/${NPU_DEVICE}"
if [ ! -e "$NPU_DEVICE_PATH" ]; then
    echo "错误: NPU设备 $NPU_DEVICE_PATH 不存在"
    echo "可用的NPU设备:"
    ls -1 /dev/davinci* 2>/dev/null || echo "  未找到任何NPU设备"
    exit 1
fi

# 只挂载指定的NPU设备和管理设备
NPU_DEVICE_ARGS="--device=${NPU_DEVICE_PATH}"

# 挂载必要的管理设备
if [ -e "/dev/davinci_manager" ]; then
    NPU_DEVICE_ARGS="$NPU_DEVICE_ARGS --device=/dev/davinci_manager"
fi
if [ -e "/dev/devmm_svm" ]; then
    NPU_DEVICE_ARGS="$NPU_DEVICE_ARGS --device=/dev/devmm_svm"
fi
if [ -e "/dev/hisi_hdc" ]; then
    NPU_DEVICE_ARGS="$NPU_DEVICE_ARGS --device=/dev/hisi_hdc"
fi

echo "挂载的NPU设备: ${NPU_DEVICE} (设备ID: ${NPU_DEVICE_ID})"

docker run -d \
  --name ${CONTAINER_NAME} \
  --restart=unless-stopped \
  --privileged \
  -p ${HOST_PORT}:${CONTAINER_PORT} \
  $NPU_DEVICE_ARGS \
  -e ASCEND_VISIBLE_DEVICES=${NPU_DEVICE_ID} \
  -e ASCEND_RT_VISIBLE_DEVICES=${NPU_DEVICE_ID} \
  -e ASCEND_DEVICE_ID=${NPU_DEVICE_ID} \
  -e MODEL_PATH=${MODEL_PATH} \
  -e USE_BF16=true \
  -e PORT=${CONTAINER_PORT} \
  -e ENABLE_CPU_OFFLOAD=false \
  -e COMPILE_MODEL=false \
  -e ATTENTION_BACKEND=sdpa \
  -v /usr/local/dcmi:/usr/local/dcmi \
  -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
  -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
  -v /etc/ascend_install.info:/etc/ascend_install.info \
  -v /data0:/data0:ro \
  ${IMAGE_NAME}

# 等待服务启动
echo ""
echo "等待服务启动..."
sleep 5

# 检查服务状态
echo ""
echo "=========================================="
echo "部署完成！"
echo "=========================================="
echo ""
echo "服务信息:"
echo "  容器名称: ${CONTAINER_NAME}"
echo "  NPU设备:  ${NPU_DEVICE}"
echo "  端口:     ${HOST_PORT}"
echo "  模型路径: ${MODEL_PATH}"
echo ""
echo "服务状态:"
docker ps | grep ${CONTAINER_NAME} || echo "警告: 容器可能未成功启动"
echo ""
echo "API地址:"
echo "  健康检查:    http://localhost:${HOST_PORT}/health"
echo "  API文档:     http://localhost:${HOST_PORT}/docs"
echo "  图像生成:    http://localhost:${HOST_PORT}/v1/images/generations"
echo "  简化接口:    http://localhost:${HOST_PORT}/v1/generate"
echo "  批量生成:    http://localhost:${HOST_PORT}/v1/generate/batch"
echo ""
echo "测试命令:"
echo "  curl http://localhost:${HOST_PORT}/health"
echo ""
echo "示例生成图像:"
echo "  curl -X POST http://localhost:${HOST_PORT}/v1/generate \\"
echo "    -F \"prompt=A beautiful sunset over the ocean\" \\"
echo "    -F \"height=1024\" \\"
echo "    -F \"width=1024\" \\"
echo "    --output test_image.png"
echo ""
echo "查看日志:"
echo "  docker logs -f ${CONTAINER_NAME}"
echo ""
echo "查看NPU使用情况:"
echo "  npu-smi info"

